/**
 * 
 */
package cs534.crf;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import org.apache.commons.math3.random.MersenneTwister;

import cs534.EmpiricalDistribution;
import cs534.crf.ConditionalRandomField.Clique;
import cs534.policy.BoltzmannPolicy;
import cs534.policy.Function2;

/**
 * @author jhostetler
 *
 */
public class CrfBoltzmannPolicy<S> extends BoltzmannPolicy<CrfBoltzmannPolicy<S>, S, int[]>
{
	protected final ConditionalRandomField<S> crf_;
	protected final CrfSampler<S> sampler_;
	protected final int Nsamples_;
	private EmpiricalDistribution<int[]> pi_ = null;
	
	// TODO: What about seeding?
	private final MersenneTwister rng_ = new MersenneTwister();
	
	/**
	 * @param parameters
	 * @param features
	 */
	public CrfBoltzmannPolicy( final double[] parameters,
							   final ArrayList<Function2<Double, S, int[]>> features,
							   final ConditionalRandomField<S> crf, 
							   final CrfSampler<S> sampler, final int Nsamples )
	{
		super( parameters, features );
		crf_ = crf;
		sampler_ = sampler;
		Nsamples_ = Nsamples;
	}
	
	/**
	 * @param parameters
	 * @param features
	 */
	public CrfBoltzmannPolicy( final double[] parameters,
							   final ArrayList<Function2<Double, S, int[]>> features,
							   final ConditionalRandomField<S> crf, 
							   final CrfSampler<S> sampler, final int Nsamples,
							   final Function2<int[], S, int[]> feature_filter )
	{
		super( parameters, features, feature_filter );
		crf_ = crf;
		sampler_ = sampler;
		Nsamples_ = Nsamples;
	}
	
	@Override
	public void enterObservationImpl( final S s )
	{
		sampler_.enterObservation( s );
		pi_ = new EmpiricalDistribution<int[]>( sampler_, Nsamples_ );
	}

	/**
	 * @see cs534.policy.BoltzmannPolicy#marginal_pi(int)
	 */
	@Override
	public double marginal_pi( final int feature )
	{
		double sum = 0.0;
		for( Map.Entry<int[], double[]> e : pi_.sparse_distribution().entrySet() ) {
			sum += e.getValue()[0] * features_.get( feature ).apply( s_, e.getKey() );
		}
		return sum;
	}

	@Override
	public CrfBoltzmannPolicy<S> update( final double[] parameters )
	{
		final List<ConditionalRandomField.Clique<S>> new_cliques 
			= new ArrayList<ConditionalRandomField.Clique<S>>();
		int param = 0;
		for( final ConditionalRandomField.Clique<S> clique : crf_.cliques() ) {
			new_cliques.add( clique.update( Arrays.copyOfRange( parameters, param, param + clique.numParameters() ) ) );
			param += clique.numParameters();
		}
		assert( param == parameters.length );
		
		final Clique<S>[] array_type = new Clique[0];
		final ConditionalRandomField<S> updated_crf = new ConditionalRandomField<S>( 
			crf_.cardinalities(), crf_.adjacency_list(), new_cliques.toArray( array_type ) );
		
		return new CrfBoltzmannPolicy<S>( 
			parameters, features_, updated_crf, sampler_.update( updated_crf ), Nsamples_, feature_filter_ );
	}

	@Override
	public int[] getAction()
	{
		double r = rng_.nextDouble();
		int count = pi_.sparse_distribution().size() - 1;
		for( Map.Entry<int[], double[]> e : pi_.sparse_distribution().entrySet() ) {
			r -= e.getValue()[0];
			if( r <= 0 || count-- == 0 ) {
				return e.getKey();
			}
		}
		throw new AssertionError( "Unreachable code" );
	}
	
	public ConditionalRandomField<S> crf()
	{
		return crf_;
	}

}
